import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np  # 显式导入 numpy
from models.ncf import NCF


# 数据集类
class RatingDataset(Dataset):
    def __init__(self, df):
        self.users = df['user_id'].values
        self.items = df['course_id'].values
        self.ratings = df['rating'].values.astype(np.float32)  # ⭐ 强制转换为 float32

    def __len__(self):
        return len(self.users)

    def __getitem__(self, idx):
        return {
            'user': self.users[idx],
            'item': self.items[idx],
            'rating': self.ratings[idx] / 5.0  # 归一化到 [0,1]
        }

# 训练函数
def train():
    global loss
    df = pd.read_csv("data/ratings.csv")

    # 修改计算方式，使用最大值而不是唯一值数量
    num_users = df['user_id'].max() + 1  # 使用最大ID + 1
    num_items = df['course_id'].max() + 1  # 使用最大ID + 1

    dataset = RatingDataset(df)
    loader = DataLoader(dataset, batch_size=32, shuffle=True)

    # 尝试不同的MLP层配置
    model = NCF(num_users, num_items, mlp_layers=[32, 16, 8])
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(10):
        for batch in loader:
            users = torch.LongTensor(batch['user'])
            items = torch.LongTensor(batch['item'])
            ratings = torch.FloatTensor(batch['rating'])
            
            # 添加调试信息
            print(f"Users shape: {users.shape}, dtype: {users.dtype}")
            print(f"Items shape: {items.shape}, dtype: {items.dtype}")
            
            preds = model(users, items)
            loss = criterion(preds, ratings)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

    torch.save(model.state_dict(), "ncf_model.pth")
    print("Model saved.")

if __name__ == "__main__":
    train()